import json
import pickle
import uuid
from abc import ABC, abstractmethod

import pandas as pd
import redis


class Cache(ABC):
    @abstractmethod
    def generate_id(self, *args, **kwargs):
        pass

    @abstractmethod
    def get(self, id, field):
        pass

    @abstractmethod
    def get_all(self, field_list) -> list:
        pass

    @abstractmethod
    def set(self, id, field, value):
        pass

    @abstractmethod
    def delete(self, id):
        pass


class RedisCache(Cache):
    def __init__(self, redis_url="redis://:123456@122.51.212.243:6379/0"):
        self.redis_client = redis.from_url(redis_url)

    def generate_id(self, *args, **kwargs):
        return str(uuid.uuid4())

    def _make_key(self, id, field):
        return f"text2sql:{id}:{field}"

    def _serialize_value(self, value):
        """序列化值，处理不同的数据类型"""
        if isinstance(value, pd.DataFrame):
            return pickle.dumps(value)
        elif isinstance(value, (str, int, float, bool)):
            return json.dumps(value)
        else:
            return pickle.dumps(value)

    def _deserialize_value(self, value):
        """反序列化值，处理不同的数据类型"""
        if value is None:
            return None
        
        try:
            # 尝试JSON反序列化
            return json.loads(value)
        except:
            try:
                # 如果JSON反序列化失败，尝试pickle反序列化
                return pickle.loads(value)
            except:
                # 如果都失败了，返回原始值
                return value

    def set(self, id, field, value):
        """设置缓存值，确保正确序列化"""
        key = self._make_key(id, field)
        serialized_value = self._serialize_value(value)
        self.redis_client.set(key, serialized_value)

    def get(self, id, field):
        """获取缓存值，确保正确反序列化"""
        key = self._make_key(id, field)
        value = self.redis_client.get(key)
        return self._deserialize_value(value)

    def get_all(self, field_list) -> list:
        """获取所有匹配的缓存值"""
        # 获取所有以text2sql:开头的键
        keys = self.redis_client.keys("text2sql:*")
        # 提取所有唯一的ID
        ids = set()
        for key in keys:
            key_parts = key.decode().split(":")
            if len(key_parts) >= 2:
                ids.add(key_parts[1])
        
        # 构建结果
        result = []
        for id in ids:
            item = {"id": id}
            for field in field_list:
                item[field] = self.get(id=id, field=field)
            result.append(item)
        return result

    def delete(self, id):
        """删除指定ID的所有缓存值"""
        keys = self.redis_client.keys(f"text2sql:{id}:*")
        if keys:
            self.redis_client.delete(*keys)

    def get_all_questions(self):
        """获取所有问题的历史记录"""
        try:
            # 获取所有以text2sql:开头的键
            keys = self.redis_client.keys("text2sql:*")
            if not keys:
                return []

            # 提取所有唯一的ID
            question_history = []
            seen_ids = set()

            for key in keys:
                key_parts = key.decode().split(":")
                if len(key_parts) >= 3:
                    id = key_parts[1]
                    if id not in seen_ids:
                        seen_ids.add(id)
                        # 获取这个ID的问题和SQL
                        question = self.get(id=id, field='question')
                        sql = self.get(id=id, field='sql')
                        if question and sql:
                            question_history.append({
                                "id": id,
                                "question": question,
                                "sql": sql
                            })

            # 按照最新的记录在前面的顺序返回
            return list(reversed(question_history))
        except Exception as e:
            print(f"Error in get_all_questions: {str(e)}")
            return []
class MemoryCache(Cache):
    def __init__(self):
        self.cache = {}

    def generate_id(self, *args, **kwargs):
        return str(uuid.uuid4())

    def set(self, id, field, value):
        if id not in self.cache:
            self.cache[id] = {}

        self.cache[id][field] = value

    def get(self, id, field):
        if id not in self.cache:
            return None

        if field not in self.cache[id]:
            return None

        return self.cache[id][field]

    def get_all(self, field_list) -> list:
        return [
            {
                "id": id,
                **{
                    field: self.get(id=id, field=field)
                    for field in field_list
                }
            }
            for id in self.cache
        ]

    def delete(self, id):
        if id in self.cache:
            del self.cache[id]